# 手写数字识别Web应用（强化字典处理）
import pickle
import numpy as np
import gradio as gr
from PIL import Image

# 加载KNN模型
def load_knn_model(model_path):
    try:
        with open(model_path, 'rb') as f:
            return pickle.load(f)
    except FileNotFoundError:
        print(f"错误：未找到模型文件 {model_path}")
        return None
    except Exception as e:
        print(f"模型加载错误：{str(e)}")
        return None

# 加载最优模型
best_knn = load_knn_model('best_knn_model.pkl')

# 强化的字典解析函数
def extract_image_from_dict(dict_data):
    """从复杂字典中提取图像数据，尝试多种可能的解析路径"""
    # 尝试直接获取常见的图像数据键
    for key in ['image', 'data', 'mask', 'pixels', 'array']:
        if key in dict_data:
            value = dict_data[key]
            # 如果值仍是字典，递归解析
            if isinstance(value, dict):
                return extract_image_from_dict(value)
            # 如果是列表或数组，直接返回
            if isinstance(value, (list, np.ndarray)):
                return value
    
    # 尝试解析可能的嵌套结构
    for key, value in dict_data.items():
        if isinstance(value, dict):
            result = extract_image_from_dict(value)
            if result is not None:
                return result
        # 检查是否包含"shape"和"data"的典型图像结构
        if key == 'shape' and 'data' in dict_data:
            return np.array(dict_data['data']).reshape(dict_data['shape'])
    
    return None

# 数据格式转换函数（增强版）
def convert_to_valid_image(drawn_data):
    """将各种格式数据转换为标准numpy图像数组"""
    # 处理字典格式（重点强化）
    if isinstance(drawn_data, dict):
        drawn_data = extract_image_from_dict(drawn_data)
        if drawn_data is None:
            raise ValueError("无法从字典中提取图像数据")
    
    # 处理列表格式
    if isinstance(drawn_data, list):
        drawn_data = np.array(drawn_data)
    
    # 确保是numpy数组
    if not isinstance(drawn_data, np.ndarray):
        raise ValueError(f"无法处理的数据类型: {type(drawn_data)}")
    
    # 处理维度问题
    if len(drawn_data.shape) == 2:
        drawn_data = np.expand_dims(drawn_data, axis=-1)
    
    # 处理通道数问题
    if drawn_data.shape[-1] == 1:
        drawn_data = np.repeat(drawn_data, 3, axis=-1)
    elif drawn_data.shape[-1] == 4:
        drawn_data = drawn_data[:, :, :3]
    
    # 确保数据类型是uint8
    if drawn_data.dtype != np.uint8:
        # 特殊处理浮点数（如果在0-1范围内）
        if np.issubdtype(drawn_data.dtype, np.floating):
            drawn_data = (drawn_data * 255).astype(np.uint8)
        else:
            drawn_data = drawn_data.astype(np.uint8)
    
    return drawn_data

# 预测函数
def predict_digit(drawn_data):
    if best_knn is None:
        return "模型加载失败，请检查模型文件"
    
    if drawn_data is None:
        return "请在画布上手写绘制一个数字（0-9）"
    
    try:
        # 数据格式转换
        try:
            image_array = convert_to_valid_image(drawn_data)
        except Exception as e:
            return f"数据转换错误: {str(e)}"
        
        # 图像处理
        img = Image.fromarray(image_array).convert('L')
        img_resized = img.resize((8, 8), Image.LANCZOS)
        img_array = np.array(img_resized)
        img_inverted = 255 - img_array
        img_normalized = (img_inverted / 255.0) * 16
        img_flat = img_normalized.flatten().reshape(1, -1)
        
        # 预测
        prediction = best_knn.predict(img_flat)
        confidence = best_knn.predict_proba(img_flat).max() * 100
        
        return f"预测结果: {prediction[0]}\n置信度: {confidence:.2f}%"
    
    except Exception as e:
        return f"预测错误: {str(e)}"

# 创建界面
def create_interface():
    with gr.Blocks(title="手写数字识别") as interface:
        gr.Markdown("# 手写数字识别器")
        gr.Markdown("""
        ### 操作步骤：
        1. 点击下方区域，选择"Draw"模式
        2. 按住鼠标左键绘制数字（0-9）
        3. 点击"预测数字"查看结果
        """)
        
        with gr.Row():
            canvas = gr.Sketchpad(label="手写区域")
            result = gr.Textbox(label="预测结果", lines=3)
        
        with gr.Row():
            predict_btn = gr.Button("预测数字", variant="primary")
            clear_btn = gr.Button("清除画布")
        
        # 绑定事件
        predict_btn.click(
            fn=predict_digit,
            inputs=canvas,
            outputs=result
        )
        clear_btn.click(
            fn=lambda: None,
            inputs=None,
            outputs=canvas
        )
    
    return interface

# 启动应用
if __name__ == "__main__":
    app = create_interface()
    app.launch(debug=True)
